import os
import argparse
import numpy as np
import scipy.io as sio
from tqdm import tqdm
from scipy import signal
from scipy.interpolate import interp1d
import random

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--distortion_level", type=str, default="moderate", 
                       choices=["moderate"], help="失真程度 - 中等级别，目标测试准确率80%左右")
    parser.add_argument("--noise_type", type=str, default="mixed", 
                       choices=["white", "pink", "brown", "mixed"], help="噪声类型")
    return parser.parse_args()

def generate_pink_noise(N, alpha=1.0):
    """生成粉色噪声 (1/f噪声)"""
    white_noise = np.random.randn(N)
    X = np.fft.fft(white_noise)
    
    freqs = np.fft.fftfreq(N)
    freqs[0] = 1e-10
    
    S = 1 / (np.abs(freqs) ** (alpha/2))
    S[0] = S[1]
    S = np.minimum(S, 50)  # 降低增益限制
    
    pink_noise = np.real(np.fft.ifft(X * S))
    return pink_noise / np.std(pink_noise)

def generate_brown_noise(N):
    """生成褐色噪声 (1/f²噪声)"""
    brown_noise = np.zeros(N)
    brown_noise[0] = np.random.randn()
    
    for i in range(1, N):
        brown_noise[i] = 0.99 * brown_noise[i-1] + np.random.randn()
    
    return brown_noise / np.std(brown_noise)

def apply_frequency_domain_distortion(signal_data, distortion_level):
    """应用频域失真 - 中等强度"""
    fft_data = np.fft.fft(signal_data)
    freq_mask = np.ones_like(fft_data)
    
    # 中等强度失真
    suppress_factor_range = (0.4, 0.7)   # 中等抑制
    enhance_factor_range = (1.5, 2.5)    # 中等增强
    noise_level = 0.15                   # 中等噪声水平
    
    # 抑制2-3个频段
    for _ in range(np.random.randint(2, 4)):
        suppress_start = int(len(fft_data) * np.random.uniform(0.1, 0.6))
        suppress_width = int(len(fft_data) * np.random.uniform(0.1, 0.3))
        suppress_end = min(suppress_start + suppress_width, len(fft_data))
        
        suppress_factor = np.random.uniform(*suppress_factor_range)
        freq_mask[suppress_start:suppress_end] = suppress_factor
        
        if suppress_start > len(fft_data)//2:
            mirror_end = len(fft_data) - suppress_start + 1
            mirror_start = len(fft_data) - suppress_end + 1
            freq_mask[mirror_start:mirror_end] = suppress_factor
    
    # 增强2-3个频段
    for _ in range(np.random.randint(2, 3)):
        enhance_start = int(len(fft_data) * np.random.uniform(0.1, 0.7))
        enhance_width = np.random.randint(30, 70)
        enhance_end = min(enhance_start + enhance_width, len(fft_data))
        
        enhance_factor = np.random.uniform(*enhance_factor_range)
        freq_mask[enhance_start:enhance_end] = enhance_factor
        
        if enhance_start > len(fft_data)//2:
            mirror_end = len(fft_data) - enhance_start + 1
            mirror_start = len(fft_data) - enhance_end + 1
            freq_mask[mirror_start:mirror_end] = enhance_factor
    
    # 添加中等强度噪声
    full_freq_noise = noise_level * np.std(np.abs(fft_data)) * \
                      (np.random.randn(len(fft_data)) + 
                       1j * np.random.randn(len(fft_data)))
    
    fft_data = fft_data * freq_mask + full_freq_noise
    
    if len(fft_data) % 2 == 0:
        fft_data[len(fft_data)//2+1:] = np.conj(np.flip(fft_data[1:len(fft_data)//2]))
    else:
        fft_data[(len(fft_data)+1)//2:] = np.conj(np.flip(fft_data[1:(len(fft_data)+1)//2]))
    
    return np.real(np.fft.ifft(fft_data))

def add_mechanical_distortions(signal_data, distortion_level):
    """添加机械失真 - 中等强度"""
    t = np.arange(len(signal_data))
    distorted_signal = signal_data.copy()
    
    # 中等强度机械失真参数
    step_freq_range = (0.02, 0.08)
    step_amp_factor = 0.15  # 中等步进幅度
    drift_amp = 0.1         # 中等漂移幅度
    pattern_repeats = (2, 4)
    
    # 1. 中等步进信号
    step_freq = np.random.uniform(*step_freq_range)
    step_amplitude = step_amp_factor * np.std(signal_data)
    step_signal = step_amplitude * np.sign(np.sin(2 * np.pi * step_freq * t))
    phase_noise = 0.3 * np.random.randn(len(t))
    step_signal *= (1 + 0.2 * np.sin(phase_noise))
    distorted_signal += step_signal
    
    # 2. 中等频率漂移
    drift_factor = 1 + 0.01 * np.sin(2 * np.pi * t / len(t)) + \
                   0.008 * np.sin(4 * np.pi * t / len(t)) + \
                   0.005 * np.sin(6 * np.pi * t / len(t))
    
    t_original = np.arange(len(distorted_signal))
    t_drifted = np.cumsum(drift_factor)
    t_drifted = t_drifted * (len(t_original) / t_drifted[-1])
    
    interp_func = interp1d(t_original, distorted_signal, kind='cubic', 
                          bounds_error=False, fill_value='extrapolate')
    distorted_signal = interp_func(t_drifted)
    
    # 3. 中等重复模式
    for _ in range(np.random.randint(2, 3)):
        pattern_length = np.random.randint(40, 80)
        pattern_start = np.random.randint(0, len(distorted_signal) - pattern_length)
        pattern = distorted_signal[pattern_start:pattern_start + pattern_length]
        
        num_repeats = np.random.randint(*pattern_repeats)
        for _ in range(num_repeats):
            insert_pos = np.random.randint(0, len(distorted_signal) - pattern_length)
            blend_factor = 0.7  # 70%混合
            distorted_signal[insert_pos:insert_pos + pattern_length] = \
                blend_factor * pattern + (1 - blend_factor) * \
                distorted_signal[insert_pos:insert_pos + pattern_length]
    
    return distorted_signal

def add_burst_noise(signal_data, distortion_level):
    """添加突发噪声 - 中等强度"""
    num_bursts = min(4, max(2, len(signal_data) // 1000))
    burst_amp_factor = 1.5
    max_burst_length = 10
    
    distorted_signal = signal_data.copy()
    
    for _ in range(num_bursts):
        burst_loc = np.random.randint(0, len(signal_data) - max_burst_length)
        burst_length = np.random.randint(3, max_burst_length)
        
        x = np.linspace(-2, 2, burst_length)
        burst_amp = burst_amp_factor * np.std(signal_data)
        
        burst_type = np.random.randint(0, 3)
        if burst_type == 0:
            burst_shape = burst_amp * np.exp(-x**2)
        elif burst_type == 1:
            burst_shape = burst_amp * np.sin(np.pi * x) * np.exp(-abs(x))
        else:
            burst_shape = burst_amp * np.random.uniform(-1, 1, burst_length)
        
        distorted_signal[burst_loc:burst_loc + burst_length] += burst_shape
    
    return distorted_signal

def add_data_loss(signal_data, distortion_level):
    """添加数据丢失 - 中等强度"""
    loss_rate = 0.01  # 1%数据丢失
    
    missing_count = max(6, int(loss_rate * len(signal_data)))
    distorted_signal = signal_data.copy()
    
    for _ in range(missing_count//2):
        start_pos = np.random.randint(0, len(signal_data) - 8)
        length = np.random.randint(2, 8)
        end_pos = min(start_pos + length, len(signal_data))
        
        if start_pos > 0 and end_pos < len(signal_data):
            lin_interp = np.linspace(distorted_signal[start_pos-1], 
                                   distorted_signal[end_pos], 
                                   end_pos - start_pos + 2)[1:-1]
            interp_noise = 0.05 * np.std(distorted_signal) * np.random.randn(len(lin_interp))
            distorted_signal[start_pos:end_pos] = lin_interp + interp_noise
    
    return distorted_signal

def add_nonlinear_distortion(signal_data, distortion_level):
    """添加非线性失真 - 中等强度"""
    quad_factor = 0.04
    cubic_factor = 0.03
    
    saturation_threshold = 2.0 * np.std(signal_data)
    
    quad_term = quad_factor * signal_data ** 2
    cubic_term = cubic_factor * signal_data ** 3
    
    distorted_signal = signal_data + quad_term - cubic_term
    
    distorted_signal = np.tanh(distorted_signal / saturation_threshold) * saturation_threshold
    
    return distorted_signal

def add_quantization(signal_data, distortion_level):
    """添加量化失真 - 中等强度"""
    bits = np.random.randint(8, 10)  # 8-9位精度
    
    signal_range = np.max(signal_data) - np.min(signal_data)
    quant_levels = 2 ** bits
    quant_step = signal_range / quant_levels
    
    quant_noise = np.random.uniform(-quant_step/2, quant_step/2, len(signal_data))
    
    return np.round(signal_data / quant_step) * quant_step + quant_noise

def add_power_interference(signal_data, distortion_level):
    """添加电源干扰 - 中等强度"""
    t = np.arange(len(signal_data))
    
    power_freq = 50 + np.random.uniform(-2, 2)  # 48-52Hz
    interference_level = 0.05  # 中等干扰
    
    total_interference = 0
    for harmonic in [1, 2]:
        harmonic_level = interference_level / harmonic
        harmonic_freq = power_freq * harmonic
        total_interference += harmonic_level * np.std(signal_data) * \
                            np.sin(2 * np.pi * harmonic_freq * t / 1000)
    
    return signal_data + total_interference

def apply_resonance_filter(signal_data, distortion_level):
    """应用共振滤波器 - 中等强度"""
    resonance_freq = np.random.uniform(0.15, 0.25)
    q_factor = np.random.uniform(5, 10)  # 中等Q值
    
    nyquist = 0.5
    norm_freq = resonance_freq / nyquist
    
    b, a = signal.butter(3, [norm_freq * 0.9, norm_freq * 1.1], btype='band')
    
    filtered_signal = signal.filtfilt(b, a, signal_data)
    
    mix_factor = 0.3  # 30%滤波信号
    return (1 - mix_factor) * signal_data + mix_factor * filtered_signal

def add_random_spikes(signal_data):
    """添加随机尖峰 - 中等强度"""
    num_spikes = np.random.randint(3, 8)
    spike_amplitude = 2.0 * np.std(signal_data)
    
    distorted_signal = signal_data.copy()
    
    for _ in range(num_spikes):
        spike_pos = np.random.randint(0, len(signal_data))
        spike_width = np.random.randint(1, 3)
        
        for j in range(spike_width):
            if spike_pos + j < len(signal_data):
                distorted_signal[spike_pos + j] += spike_amplitude * np.random.uniform(-1, 1)
    
    return distorted_signal

def add_low_frequency_drift(signal_data):
    """添加低频漂移 - 中等强度"""
    t = np.arange(len(signal_data))
    t_norm = t / len(t)
    
    drift = 0.1 * np.std(signal_data) * (
        np.sin(2 * np.pi * 0.3 * t_norm) +
        0.5 * np.sin(2 * np.pi * 0.2 * t_norm)
    )
    
    return signal_data + drift

def generate_complex_noise(signal_data, noise_type, noise_level):
    """生成复杂噪声 - 中等强度"""
    if noise_type == "white":
        noise = np.random.randn(len(signal_data))
    elif noise_type == "pink":
        noise = generate_pink_noise(len(signal_data))
    elif noise_type == "brown":
        noise = generate_brown_noise(len(signal_data))
    else:  # mixed
        white = np.random.randn(len(signal_data))
        pink = generate_pink_noise(len(signal_data))
        brown = generate_brown_noise(len(signal_data))
        
        weights = np.random.dirichlet([1, 1, 1])
        noise = weights[0] * white + weights[1] * pink + weights[2] * brown
    
    noise = noise / np.std(noise) * noise_level * np.std(signal_data)
    return signal_data + noise

def simulate_realistic_distortion(signal_data, distortion_level, noise_type):
    """应用综合失真处理 - 中等强度版本"""
    if signal_data.ndim == 1:
        signal_data = signal_data.reshape(-1, 1).flatten()
    
    # 中等噪声水平
    noise_level = 0.12  # 12%噪声
    
    # 随机选择要应用的失真类型（70-80%概率）
    apply_freq_dist = np.random.random() < 0.8
    apply_mechanical = np.random.random() < 0.7
    apply_quantization = np.random.random() < 0.8
    apply_resonance = np.random.random() < 0.7
    apply_data_loss = np.random.random() < 0.7
    apply_burst = np.random.random() < 0.7
    apply_nonlinear = np.random.random() < 0.8
    apply_power = np.random.random() < 0.7
    apply_spikes = np.random.random() < 0.7
    apply_lf_drift = np.random.random() < 0.7
    
    distorted_signal = signal_data.copy()
    
    # 1. 频域失真
    if apply_freq_dist:
        distorted_signal = apply_frequency_domain_distortion(distorted_signal, distortion_level)
    
    # 2. 添加复杂噪声
    distorted_signal = generate_complex_noise(distorted_signal, noise_type, noise_level)
    
    # 3. 机械失真
    if apply_mechanical:
        distorted_signal = add_mechanical_distortions(distorted_signal, distortion_level)
    
    # 4. 低频漂移
    if apply_lf_drift:
        distorted_signal = add_low_frequency_drift(distorted_signal)
    
    # 5. 量化失真
    if apply_quantization:
        distorted_signal = add_quantization(distorted_signal, distortion_level)
    
    # 6. 共振滤波
    if apply_resonance:
        distorted_signal = apply_resonance_filter(distorted_signal, distortion_level)
    
    # 7. 数据丢失
    if apply_data_loss:
        distorted_signal = add_data_loss(distorted_signal, distortion_level)
    
    # 8. 突发噪声
    if apply_burst:
        distorted_signal = add_burst_noise(distorted_signal, distortion_level)
    
    # 9. 随机尖峰
    if apply_spikes:
        distorted_signal = add_random_spikes(distorted_signal)
    
    # 10. 非线性失真
    if apply_nonlinear:
        distorted_signal = add_nonlinear_distortion(distorted_signal, distortion_level)
    
    # 11. 电源干扰
    if apply_power:
        distorted_signal = add_power_interference(distorted_signal, distortion_level)
    
    # 12. 添加第二轮复杂噪声
    extra_noise_level = 0.05
    distorted_signal = generate_complex_noise(distorted_signal, "mixed", extra_noise_level)
    
    # 13. 最终缩放
    final_scale = np.random.uniform(0.8, 1.2)  # ±20%缩放
    distorted_signal *= final_scale
    
    # 14. 添加中等强度的周期性故障特征扰动
    t = np.arange(len(distorted_signal))
    periodic_distortion = 0.05 * np.std(distorted_signal) * (
        np.sin(2 * np.pi * t / 50) + 
        0.5 * np.sin(2 * np.pi * t / 70)
    )
    distorted_signal += periodic_distortion
    
    # 15. 添加中等密度的高频毛刺
    spike_density = 0.01  # 1%的点有毛刺
    num_spikes = int(spike_density * len(distorted_signal))
    spike_positions = np.random.choice(len(distorted_signal), num_spikes, replace=False)
    for pos in spike_positions:
        distorted_signal[pos] += np.random.uniform(-1, 1) * np.std(distorted_signal) * 0.3
    
    # 16. 最终随机扰动
    final_noise = 0.03 * np.std(distorted_signal) * np.random.randn(len(distorted_signal))
    distorted_signal += final_noise
    
    # 17. 添加中等强度的阶跃失真
    if np.random.random() > 0.5:  # 50%概率
        step_position = np.random.randint(len(distorted_signal) // 4, 3 * len(distorted_signal) // 4)
        step_amplitude = np.random.uniform(0.3, 0.7) * np.std(distorted_signal)
        distorted_signal[step_position:] += step_amplitude * np.random.choice([-1, 1])
    
    return distorted_signal

def load_cwru_data(data_path):
    """加载训练集数据"""
    train_path = os.path.join(data_path, "train")
    if not os.path.exists(train_path):
        raise FileNotFoundError(f"训练数据目录 {train_path} 不存在，请先运行 preprocess.py")
    
    data_files = [f for f in os.listdir(train_path) if f.endswith('.npy')]
    processed_data = {}
    
    for file in tqdm(data_files, desc="加载训练数据文件"):
        fault_type = file.replace('.npy', '')
        file_path = os.path.join(train_path, file)
        segments = np.load(file_path)
        processed_data[fault_type] = segments
        print(f"加载 {fault_type} 训练数据，形状: {segments.shape}")
    
    return processed_data

def generate_simulated_data(original_data, args):
    """生成仿真数据"""
    simulated_data = {}
    
    for fault_type, segments in tqdm(original_data.items(), desc="生成仿真数据"):
        simulated_segments = []
        
        for segment in segments:
            # 应用综合失真
            simulated_segment = simulate_realistic_distortion(
                segment, args.distortion_level, args.noise_type)
            simulated_segments.append(simulated_segment)
        
        simulated_data[fault_type] = np.array(simulated_segments)
    
    return simulated_data

def save_processed_data(processed_data, save_path):
    """保存处理后的数据"""
    os.makedirs(save_path, exist_ok=True)
    
    for fault_type, segments in processed_data.items():
        data_array = np.array(segments)
        save_file = os.path.join(save_path, f"{fault_type}.npy")
        np.save(save_file, data_array)
        print(f"已保存 {fault_type} 类别的数据，形状: {data_array.shape}")

def main():
    args = parse_args()
    
    # 设置路径
    data_path = "./cwru_prepro"
    save_path = "./cwru_sim_real"
    
    # 加载训练集数据
    print("加载CWRU训练集数据...")
    original_data = load_cwru_data(data_path)
    
    # 生成仿真数据
    print(f"基于训练集生成仿真数据（失真程度：{args.distortion_level}，噪声类型：{args.noise_type}）...")
    simulated_data = generate_simulated_data(original_data, args)
    
    # 保存仿真数据
    print("保存仿真数据...")
    save_processed_data(simulated_data, save_path)
    print("仿真数据生成完成！")

if __name__ == "__main__":
    main() 